import torch.cuda
from torch import nn
from torchmetrics.functional import pairwise_cosine_similarity

from common_models.MLP import MLP

device = 'cuda' if torch.cuda.is_available() else 'cpu'


class AttrEmbedModel(nn.Module):
    def __init__(self, args,feat_dim,attr_dim):
        super(AttrEmbedModel, self).__init__()
        # 输入维度为图片特征维度，输出维度为属性的数量
        self.attr_classifier = MLP(feat_dim, attr_dim, relu=False).to(device)
        self.label_emb = MLP(feat_dim, 512, relu=True,dropout=args.dropout, norm=args.norm).to(device)

    def forward(self, imgs, labels=None, attrs=None):
        # 根据图片特征学习对象的属性
        attr_preds = self.attr_classifier(imgs)  # 属性评分
        loss_fn = torch.nn.MSELoss()
        loss = 1- torch.cosine_similarity(attr_preds, attrs,dim=1)  # N*C （no. f classes)
        # loss =loss_fn(attr_preds,attrs)
        return loss.mean(),attr_preds
